#include "npdm_patterns.h"
#include "npdm_sweep.h"
#include "npdm_operators.h"

using namespace Npdm;

//===========================================================================================================================================================

Npdm_sweep::Npdm_sweep(int order, int nsites, Npdm_patterns & npdm_patterns) 
{ 
  nsites_ = nsites; 
  order_ = order, 
  npdm_patterns_ = npdm_patterns; 
  repeated_elements_ = 0;
}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

void Npdm_sweep::build_npdm_element( std::vector< std::pair<CD,int> > & lhs,
                                     std::vector< std::pair<CD,int> > & dot,
                                     std::vector< std::pair<CD,int> > & rhs )
{

  if ( lhs.size() + dot.size() + rhs.size() != 2*order_ ) return;

  // Combine LHS, Dot, RHS
  std::vector< std::pair<CD,int> > element;

  for (auto i = lhs.begin(); i != lhs.end(); i++) {
    element.push_back(*i);
  }
  for (auto i = dot.begin(); i != dot.end(); i++) {
    element.push_back(*i);
  }
  for (auto i = rhs.begin(); i != rhs.end(); i++) {
    element.push_back(*i);
  }

  // Sort all creation ops to left
  std::sort( element.begin(), element.end() );

  // Check for redundant repetitions
  bool redundant = ( npdm_elements_.find( element ) != npdm_elements_.end() );
  if ( redundant ) {
    repeated_elements_++;
    print_opstring( element );
    assert( false );
  }
  npdm_elements_.insert( element );

}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

void Npdm_sweep::build_block_operators( int dotsite, Npdm_operators & npdm_operators )
{
  // Build LHS operators
  for (auto it = npdm_patterns_.lhs_cd_begin(); it != npdm_patterns_.lhs_cd_end(); it++) {
    std::vector<CD> tmp = *it;
    npdm_operators.build_operator_set( 'l', dotsite, tmp );
  }
  // Build Dot operators
  for (auto it = npdm_patterns_.dot_cd_begin(); it != npdm_patterns_.dot_cd_end(); it++) {
    std::vector<CD> tmp = *it;
    npdm_operators.build_operator_set( 'd', dotsite, tmp );
  }
  // Build RHS operators
  for (auto it = npdm_patterns_.rhs_cd_begin(); it != npdm_patterns_.rhs_cd_end(); it++) {
    std::vector<CD> tmp = *it;
    npdm_operators.build_operator_set( 'r', dotsite, tmp );
  }
}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

void Npdm_sweep::contract_block_operators( std::vector<CD> & lhs_cd, 
                                           std::vector<CD> & dot_cd, 
                                           std::vector<CD> & rhs_cd, 
                                           Npdm_operators & npdm_operators )
{

  std::vector< std::pair<CD,int> > null;

  // Dot only
  if ( lhs_cd.size() == 0 && rhs_cd.size() == 0 ) {
    for (auto dot = npdm_operators.dot_ops.at(dot_cd).begin() ; dot != npdm_operators.dot_ops.at(dot_cd).end() ; dot++ ) {
      std::vector< std::pair<CD,int> > dotop = *dot;
      build_npdm_element( null, dotop, null );
    }
  }
  // Dot and RHS only
  else if ( lhs_cd.size() == 0 ) {
    for (auto dot = npdm_operators.dot_ops.at(dot_cd).begin() ; dot != npdm_operators.dot_ops.at(dot_cd).end() ; dot++ ) {
      std::vector< std::pair<CD,int> > dotop = *dot;
      for (auto rhs = npdm_operators.rhs_ops.at(rhs_cd).begin() ; rhs != npdm_operators.rhs_ops.at(rhs_cd).end() ; rhs++ ) {
        std::vector< std::pair<CD,int> > rhsop = *rhs;
        build_npdm_element( null, dotop, rhsop );
      }
    }
  }
  // Dot and LHS only
  else if ( rhs_cd.size() == 0 ) {
    for (auto lhs = npdm_operators.lhs_ops.at(lhs_cd).begin() ; lhs != npdm_operators.lhs_ops.at(lhs_cd).end() ; lhs++ ) {
      std::vector< std::pair<CD,int> > lhsop = *lhs;
      for (auto dot = npdm_operators.dot_ops.at(dot_cd).begin() ; dot != npdm_operators.dot_ops.at(dot_cd).end() ; dot++ ) {
        std::vector< std::pair<CD,int> > dotop = *dot;
        build_npdm_element( lhsop, dotop, null );
      }
    }
  }
  // LHS, Dot and RHS blocks
  else {
    for (auto lhs = npdm_operators.lhs_ops.at(lhs_cd).begin() ; lhs != npdm_operators.lhs_ops.at(lhs_cd).end() ; lhs++ ) {
      std::vector< std::pair<CD,int> > lhsop = *lhs;
      for (auto dot = npdm_operators.dot_ops.at(dot_cd).begin() ; dot != npdm_operators.dot_ops.at(dot_cd).end() ; dot++ ) {
        std::vector< std::pair<CD,int> > dotop = *dot;
        for (auto rhs = npdm_operators.rhs_ops.at(rhs_cd).begin() ; rhs != npdm_operators.rhs_ops.at(rhs_cd).end() ; rhs++ ) {
          std::vector< std::pair<CD,int> > rhsop = *rhs;
          build_npdm_element( lhsop, dotop, rhsop );
        }
      }
    }
  }

}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

void Npdm_sweep::dmrg_sweep()
{

  // Loop over DMRG sweep (all sites)
std::cout << "Repeated elements built in each pattern:\n";
  for (int dotsite = 1; dotsite <= nsites_; dotsite++) {

    std::cout << "dotsite = " << dotsite << std::endl;
    // Prepare block operators for all necessary site combinations
    Npdm_operators npdm_operators( nsites_ );
    build_block_operators( dotsite, npdm_operators );

    // Loop over DMRG operator patterns (-CD-) (-CD-) (-CD-)
    for (auto it = npdm_patterns_.ldr_cd_begin(); it != npdm_patterns_.ldr_cd_end(); it++) {

      // Initialize empty creation-destruction patterns for LHS and RHS if on edge of sweep
      std::vector<CD> lhs_cd, rhs_cd;
      std::vector<CD> dot_cd = it->at('d');
      if ( dotsite > 1 ) lhs_cd = it->at('l');
      if ( dotsite < nsites_ ) rhs_cd = it->at('r');
//npdm_patterns_.print_cd_string( lhs_cd );
//npdm_patterns_.print_cd_string( dot_cd );
//npdm_patterns_.print_cd_string( rhs_cd ); std::cout << std::endl;

      // Contract all block operators associated with this DMRG pattern
      contract_block_operators( lhs_cd, dot_cd, rhs_cd, npdm_operators );

    }

  }

  std::cout << "===============================================================\n";
  std::cout << "Number of distinct spatial strings built       = "<< npdm_elements_.size() << std::endl;;
  std::cout << "Number of repeated strings built               = "<< repeated_elements_ << std::endl;;

}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

bool Npdm_sweep::included_in_dmrg_sweep( std::vector<int> & indices )
{

  // Convert to operator form 
  std::vector< std::pair<CD,int> > opstring;

  // Creation operators
  for (int i=0; i < indices.size()/2; i++) {
    opstring.push_back( std::make_pair(CREATION, indices[i]) );
  }
  // Destruction operators
  for (int i=indices.size()/2; i < indices.size(); i++) {
    opstring.push_back( std::make_pair(DESTRUCTION, indices[i]) );
  }

  // Test if built in DMRG sweep
  return ( npdm_elements_.find( opstring ) != npdm_elements_.end() );

}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

void Npdm_sweep::print_opstring( std::vector< std::pair<CD,int> > opstring )
{
  char cd;
  for (auto op = opstring.begin(); op != opstring.end(); op++) {
    cd = '.';
    if (op->first == CREATION) cd = '+';
    std::cout << "(" << op->second << cd << ")";
  }
  std::cout << std::endl;
}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

void Npdm_sweep::print_elements()
{
  std::cout << "All distinct elements built:\n";
  for (auto element = npdm_elements_.begin(); element != npdm_elements_.end(); element++) {
     print_opstring( *element );
  }
}

//===========================================================================================================================================================

